
#include <unistd.h>

#include <iostream>
#include <string>

#include <http_client.h>

namespace tc = triton::client;

#define FAIL_IF_ERR(X, MSG)                                        \
  {                                                                \
    tc::Error err = (X);                                           \
    if (!err.IsOk()) {                                             \
      std::cerr << "error: " << (MSG) << ": " << err << std::endl; \
      exit(1);                                                     \
    }                                                              \
  }

void
ValidateShapeAndDatatype(
    const std::string& name, std::shared_ptr<tc::InferResult> result)
{
  std::vector<int64_t> shape;
  FAIL_IF_ERR(result->Shape(name, &shape), "unable to get shape for " + name);
  // Validate shape
  std::cout << "shape size: " << shape.size();
  std::cout <<  " shape: " << shape[0] << std::endl;
  if ((shape.size() != 1) || (shape[0] != 512)) {
    std::cerr << "error: received incorrect shapes for " << name << std::endl;
    // exit(1);
  }
  std::string datatype;
  FAIL_IF_ERR(
      result->Datatype(name, &datatype), "unable to get datatype for " + name);
  // Validate datatype
  if (datatype.compare("FP32") != 0) {
    std::cerr << "error: received incorrect datatype for " << name << ": "
              << datatype << std::endl;
    // exit(1);
  }
}

std::vector<float> embeddingText(const std::string &text) const {
    bool verbose = false;
    std::string url("localhost:8000");
    tc::Headers http_headers;

    std::string model_name = "text_embedding";
    std::string model_version = "";

    std::unique_ptr<tc::InferenceServerHttpClient> client;
    FAIL_IF_ERR(
        tc::InferenceServerHttpClient::Create(&client, url, verbose),
        "unable to create http client");

      std::vector<std::string> input0_data(1);
  input0_data[0] = text;

  // 在这里，你就得告诉 triton 到底 shape 是什么了，模型配置里的 -1 代表不确定个数的输入，
  // 现在你不能还不确定了，你得告诉 triton 了。
  // 另外，字符串形式的时候，默认就知道字符串不定字符个数，所以 -1 不是代表不确定字符串有多少字符，
  // 而是说字符串的个数不确定。如果配置文件里写 1，那只能有一个字符串输入，不是说字符串只能有一个字符。
  std::vector<int64_t> shape{ 1 };

  // Initialize the inputs with the data.
  tc::InferInput* input0;

  FAIL_IF_ERR(
      tc::InferInput::Create(&input0, "product_reviews", shape, "BYTES"),
      "unable to get product_reviews");
  std::shared_ptr<tc::InferInput> input0_ptr;
  input0_ptr.reset(input0);
  input0_ptr->SetBinaryData(true);

  FAIL_IF_ERR(
      input0_ptr->AppendFromString(input0_data),
      "unable to set data for product_reviews");

  // Generate the outputs to be requested.
  tc::InferRequestedOutput* output0;

  FAIL_IF_ERR(
      tc::InferRequestedOutput::Create(&output0, "unnorm_text_features"),
      "unable to get unnorm_text_features");
  std::shared_ptr<tc::InferRequestedOutput> output0_ptr;
  output0_ptr.reset(output0);
  output0_ptr->SetBinaryData(true);

  // The inference settings. Will be using default for now.
  tc::InferOptions options(model_name);
  options.model_version_ = model_version;

  std::vector<tc::InferInput*> inputs = {input0_ptr.get()};
  std::vector<const tc::InferRequestedOutput*> outputs = { output0_ptr.get() };

  tc::InferResult* results;
  FAIL_IF_ERR(
      client->Infer(&results, options, inputs, outputs, http_headers),
      "unable to run model");
  std::shared_ptr<tc::InferResult> results_ptr;
  results_ptr.reset(results);

  // Validate the results...
  ValidateShapeAndDatatype("unnorm_text_features", results_ptr);

  float* output0_data_ptr;

  size_t output0_byte_size = 512 * 4; // 输出 512 个 float 所以是 512 * 4 字节

  size_t recv_output0_byte_size;

  FAIL_IF_ERR(
      results_ptr->RawData(
          "unnorm_text_features", (const uint8_t**)&output0_data_ptr,
          &recv_output0_byte_size),
      "unable to get result data for 'unnorm_text_features'");
  if (recv_output0_byte_size != output0_byte_size) {
    std::cerr << "error: received incorrect byte size for 'unnorm_text_features': "
              << recv_output0_byte_size << std::endl;
    exit(1);
  }
  std::vector<float> tokens;

  std::cout << "unnorm_text_features: ";
  for (size_t i = 0; i < 512; ++i) {
    std::cout << output0_data_ptr[i] << " ";
    tokens.push_back(output0_data_ptr[i]);
  }
  std::cout << std::endl;
  return tokens;
}